#include "RodCalculator.h"

#include "DataStructure/Common/Status.h"
#include "DataStructure/Input/Solver/DOFDirections.h"

#include "DBManager/IService/Input/Node/INodeService.h"
#include "DBManager/IService/Input/Element/IElementService.h"
#include "DBManager/IService/Input/Material/IIsotropicMaterialService.h"
#include "DBManager/IService/Input/Property/IPRODService.h"
#include "DBManager/IService/Result/Nodal/IDISPLACEMENTService.h"
#include "DBManager/IService/Result/ElementalStrain/ICONRODStrainService.h"
#include "DBManager/IService/Result/ElementalStrain/IRODStrainService.h"
#include "DBManager/IService/Result/ElementalStress/ICONRODStressService.h"
#include "DBManager/IService/Result/ElementalStress/IRODStressService.h"
#include "DBManager/IService/Result/ElementalForce/ICONRODForceService.h"
#include "DBManager/IService/Result/ElementalForce/IRODForceService.h"

#include "SGFem/Pre/Assembly/Assembly.h"
#include "SGFem/Pre/Assembly/CoordinateProcessor.h"
#include "SGFem/Pre/Assembly/ModelInfoQuery.h"
#include "../Common/ElementTools.hpp"

using SG::Algebra::Matrixd;
using SG::Algebra::Vector3D;
using SG::Algebra::PointArray;
using SG::DataStructure::Common::Status;
using SG::DataStructure::Common::Id;
using SG::DataStructure::Common::IdArray;
using SG::DataStructure::Common::Real;
using namespace SG::DataStructure::FEM;
using SG::DataStructure::Post::RODStrainData;
using SG::DataStructure::Post::RODStressData;
using SG::DataStructure::Post::CONRODForceData;
using SG::DataStructure::Post::RODForceData;
using namespace SG::DBManager;
using SG::FEMSolver::AssemblyEigen;
using SG::FEMSolver::CoordinateProcessor;
using SG::FEMSolver::Element::RodCalculator;

void RodCalculator::ComputeLinkStiffness
(
const SG::DataStructure::Common::Real            L,
const SG::DataStructure::FEM::IsotropicMaterial& material,
const SG::DataStructure::Common::Real            S,
const SG::Algebra::Matrixd&                      T,
_OUT SG::Algebra::Matrixd&                       kg
)
{
    const Real ESL{material.m_E * S / L};
    const Matrixd kl (2, 2, {ESL, -ESL, -ESL, ESL});
    kg = T.Transpose () * kl * T;
}

// Torsion
void RodCalculator::ComputeTorsionStiffness
(
const SG::DataStructure::Common::Real            L,
const SG::DataStructure::FEM::IsotropicMaterial& material,
const SG::DataStructure::Common::Real            J,
const SG::Algebra::Matrixd&                      T,
_OUT SG::Algebra::Matrixd&                       kt
)
{
    const Real GJL{material.m_G * J / L};
    const Matrixd ktl (2, 2, {GJL, -GJL, -GJL, GJL});
    kt = T.Transpose () * ktl * T;
}

void RodCalculator::ComputeTemperatureLoad
(
const SG::DataStructure::FEM::IsotropicMaterial& material,
const SG::DataStructure::Common::Real            S,
const SG::DataStructure::Common::Real            deltaT,
const SG::Algebra::Matrixd&                      T,
_OUT SG::Algebra::Matrixd&                       Pe0
)
{
    const Real& E{material.m_E};
    const Real& A{material.m_A};

    Matrixd p1 (2, 1);
    p1 (0, 0) = -E * S * A * deltaT;
    p1 (1, 0) = E * S * A * deltaT;

    Pe0 = T.Transpose () * p1;
}

void RodCalculator::ComputeLumpMass (const Real m, _OUT Matrixd& lumpedM)
{
    for (size_t i{ 0 }; i < 6; ++i)
    {
        lumpedM (i, i) = m;
    }
}

void RodCalculator::ComputeCoupleMass
(
const SG::DataStructure::Common::Real     m,
const SG::Algebra::Vector3D&              xVec,
const SG::DataStructure::Common::IdArray& dispCoordIdSet,
const SG::Algebra::PointArray&            pointSet,
_OUT SG::Algebra::Matrixd&                coupledM
)
{
    const auto& coordProcessor = CoordinateProcessor::GetInstance ();
    const Matrixd T (1, 3, {xVec.m_x, xVec.m_y, xVec.m_z});
    std::vector<SG::Algebra::Matrixd> TEle2GSet (2);
    for (size_t iNode{0}; iNode < 2; ++iNode)
    {
        TEle2GSet[iNode] = T;
        if (dispCoordIdSet[iNode] > 0)
        {
            // 处理节点位移坐标系 CD
            Matrixd TG (3, 3);
            coordProcessor.GetCoordTrans (dispCoordIdSet[iNode], pointSet[iNode], _OUT TG);
            TEle2GSet[iNode] = T * TG.Transpose ();
        }
        Matrixd Mii = (SG::Algebra::identity (3) - TEle2GSet[iNode].Transpose () * TEle2GSet[iNode] * (1.0 / 6.0)) * m;
        coupledM.BlockFill (iNode * 3, iNode * 3, Mii);
    }
    const Matrixd Mab{TEle2GSet[0].Transpose () * TEle2GSet[1] * (m / 6.0)};
    coupledM.BlockFill (0, 3, Mab);
    coupledM.BlockFill (3, 0, Mab.Transpose ());
}

void RodCalculator::Initialize (SG::DataStructure::FEM::ElementTypes type)
{
    m_type = type;
}

auto GetElement = [] (const DBServiceFactorySP& dbServiceFactory, const ElementTypes& type, const std::size_t iEle, CONRODData& rodData)->Status
{
    if (ElementTypes::ELETYPE_Link == type)
    {
        CRODData tempRod;
        if (Status::STATUS_Failure == dbServiceFactory->get<IElementService> ()->FindByIndex (iEle, _OUT tempRod))
        {
            // todo: 错误日志
            return Status::STATUS_Failure;
        };

        PRODData tempPRod;
        if (Status::STATUS_Failure == dbServiceFactory->get<IPRODService> ()->FindById (tempRod.m_pId, _OUT tempPRod))
        {
            // todo: 错误日志
            return Status::STATUS_Failure;
        }
        rodData.m_id   = tempRod.m_id;
        rodData.m_g[0] = tempRod.m_g[0];
        rodData.m_g[1] = tempRod.m_g[1];
        rodData.m_A    = tempPRod.m_area;
        rodData.m_J    = tempPRod.m_torsionFactor;
        rodData.m_C    = tempPRod.m_torsionStressfactor;
        rodData.m_mId  = tempPRod.m_mId;
        rodData.m_NSM  = tempPRod.m_nsm;
    }
    else
    {
        // CONROD
        if (Status::STATUS_Failure == dbServiceFactory->get<IElementService> ()->FindByIndex (iEle, _OUT rodData))
        {
            // todo: 错误日志
            return Status::STATUS_Failure;
        }
    }

    return Status::STATUS_Success;
};

/** @brief    位移输出坐标系到单元局部坐标系下的转换矩阵 */
auto CoordTrans = [] (const Vector3D& xVec, const CoordinateProcessor& coordProcessor, const IdArray& dispCoordIds, const PointArray& points)->Matrixd
{
    Matrixd TEle (2, 6);
    const Matrixd T (1, 3, {xVec.m_x, xVec.m_y, xVec.m_z});
    for (size_t i{0}; i < 2; ++i)
    {
        if (0 == dispCoordIds[i])
        {
            TEle.BlockFill (i, i * 3, T);
        }
        else
        {
            // 处理节点位移坐标系 CD
            Matrixd TDisp (3, 3);
            coordProcessor.GetCoordTrans (dispCoordIds[i], points[i], _OUT TDisp);
            TEle.BlockFill (i, i * 3, T * TDisp.Transpose ());
        }
    }
    return TEle;
};

void RodCalculator::ComputeElementMatries (const SG::DBManager::DBServiceFactorySP dbServiceFactory)
{
    auto& assembly = AssemblyEigen::GetInstance ();
    const auto& coordProcessor = CoordinateProcessor::GetInstance ();

    Matrixd                  kg (6, 6);
    Matrixd                  kt (6, 6);
    std::vector<std::size_t> nodeIndexes (2);
    IdArray                  dispCoordIds (2);
    PointArray               points (2);

    static const DofSet DOFs{DOFDirections::DOF_Ux, DOFDirections::DOF_Uy, DOFDirections::DOF_Uz};
    static const DofSet RotDOFs{DOFDirections::DOF_Rotx, DOFDirections::DOF_Roty, DOFDirections::DOF_Rotz};

    const auto& nodeService = dbServiceFactory->get<INodeService> ();
    const auto eleSize = dbServiceFactory->get<IElementService>()->GetElementSize (m_type);
    for (std::size_t iEle{0}; iEle < eleSize; ++iEle)
    {
        // 根据不同类型，获取单元数据
        CONRODData rodData;
        GetElement (dbServiceFactory, m_type, iEle, _OUT rodData);

        // 获取节点数据：节点索引、节点全局坐标、节点位移坐标系ID
        ModelInfoQuery::QueryNode (nodeService, rodData.m_g, 2, _OUT points, _OUT nodeIndexes, _OUT dispCoordIds);

        // 单元坐标转换矩阵
        auto xVec = SG::Algebra::Distance (points[1], points[0]);
        Real L;
        xVec.Normalize (L);
        const auto& TEle = CoordTrans (xVec, coordProcessor, dispCoordIds, points);

        // 获取单元材料
        IsotropicMaterial material;
        auto status = dbServiceFactory->get<IIsotropicMaterialService> ()->FindById (rodData.m_mId, _OUT material);
        assert (Status::STATUS_Success == status);

        // 刚度矩阵计算
        std::vector<int> globalIndex;
        if (rodData.m_A > 1.0e-12)
        {
            ComputeLinkStiffness (L, material, rodData.m_A, TEle, _OUT kg);
            assembly.GetGlobalDOFId (nodeIndexes, DOFs, _OUT globalIndex);
            assembly.AccumulateStiffness (globalIndex, kg);
        }

        if (rodData.m_J > 1.0e-12)
        {
            ComputeTorsionStiffness (L, material, rodData.m_J, TEle, _OUT kt);
            assembly.GetGlobalDOFId (nodeIndexes, RotDOFs, _OUT globalIndex);
            assembly.AccumulateStiffness (globalIndex, kt);
        }

        // 质量矩阵计算
        // half of the mass of the rod, include nonstructural mass
        const Real m{0.5 * (material.m_rho * rodData.m_A + rodData.m_NSM) * L};
        Matrixd massMatrix (6, 6);
        const auto needCoupleMass = assembly.NeedCoupleMass ();
        if (needCoupleMass)
        {
            ComputeCoupleMass (m, xVec, dispCoordIds, points, _OUT massMatrix);
        }
        else
        {
            ComputeLumpMass (m, _OUT massMatrix);
        }

        std::vector<int> dofGlobalIndex;
        assembly.GetGlobalDOFId (nodeIndexes, DOFs, _OUT dofGlobalIndex);
        assembly.AccumulateCoupledMass (dofGlobalIndex, massMatrix);
    }
}

//计算轴向内力
void RodCalculator::ComputeAxialForce (const SG::Algebra::Matrixd& TEle,
                                       const Matrixd&              eleTranslationalDisplacement,
                                       const Real&                 L,
                                       const Matrixd&              coeffMatrix,
                                       const IsotropicMaterial&    material,
                                       const CONRODData&           rodData,
                                       _OUT Real&                   eleAxialStrain,
                                       _OUT Real&                   eleAxialStress,
                                       _OUT Real&                   eleAxialForce)
{
    auto    u1SubU2 = coeffMatrix * (TEle * eleTranslationalDisplacement);
    eleAxialStrain        = u1SubU2 (0, 0) / L;
    eleAxialStress = material.m_E * eleAxialStrain;
    eleAxialForce = eleAxialStress * rodData.m_A;
}

// 后处理数据
void RodCalculator::DataRecovery
(
const SG::DBManager::DBServiceFactorySP             dbServiceFactory,
int                                                 subcaseId, 
SG::DataStructure::Common::Real time,
const std::vector<Real>& globalDisplacement
)
{
    std::vector<RODStrainData>    rodStrainSet;
    std::vector<RODStressData>    rodStressSet;
    std::vector<RODForceData>     rodForceSet;

    const auto& assembly     = AssemblyEigen::GetInstance ();
    const auto& coordProcessor = CoordinateProcessor::GetInstance ();

    CONRODData               rodData;
    std::vector<std::size_t> nodeIndexes (2);
    IdArray                  dispCoordIds (2);
    PointArray               points (2);
    std::vector<int>         dofGlobalIds (6);
    Matrixd eleTranslationalDisplacement (6, 1);
    Matrixd eleRotationalDisplacement (6, 1);

    static const DofSet DOFs{DOFDirections::DOF_Ux, DOFDirections::DOF_Uy, DOFDirections::DOF_Uz};
    static const DofSet RotDOFs{DOFDirections::DOF_Rotx, DOFDirections::DOF_Roty, DOFDirections::DOF_Rotz};

    const auto& nodeService = dbServiceFactory->get<INodeService> ();
    const auto eleSize = dbServiceFactory->get<IElementService>()->GetElementSize (m_type);
    rodStrainSet.reserve(eleSize);
    rodStressSet.reserve(eleSize);
    rodForceSet.reserve(eleSize);

    const Matrixd coeffMatrix (1, 2, {-1.0, 1.0});
    for (std::size_t index{0}; index < eleSize; ++index)
    {
        // 根据不同类型，获取单元数据
        GetElement (dbServiceFactory, m_type, index, _OUT rodData);
        ModelInfoQuery::QueryNode (nodeService, rodData.m_g, 2, points, nodeIndexes, dispCoordIds);

        auto xVec = SG::Algebra::Distance (points[1], points[0]);
        Real L;
        xVec.Normalize (L);
        const auto& TEle = CoordTrans (xVec, coordProcessor, dispCoordIds, points);
        IsotropicMaterial material;
        auto status = dbServiceFactory->get<IIsotropicMaterialService> ()->FindById (rodData.m_mId, _OUT material);
        assert (Status::STATUS_Success == status);

        // 每个新单元必须重新初始化数值
        Real eleAxialStrain{0.0};
        Real eleAxialStress{0.0};
        Real eleAxialForce{0.0};
        if (rodData.m_A > 1.0e-12)
        {
            assembly.GetGlobalDOFId (nodeIndexes, DOFs, dofGlobalIds);
            ElementTools::GetElementDisplacement (dofGlobalIds, globalDisplacement, eleTranslationalDisplacement);

            //计算单元内力
            ComputeAxialForce (TEle, eleTranslationalDisplacement, L, coeffMatrix, material, rodData, eleAxialStrain, eleAxialStress, eleAxialForce);
        }

        Real eleTorsionalStrain{0.0};
        Real eleTorsionalStress{0.0};
        Real eleTorque{0.0};
        if (rodData.m_J > 1.0e-12)
        {
            assembly.GetGlobalDOFId (nodeIndexes, RotDOFs, dofGlobalIds);
            ElementTools::GetElementDisplacement (dofGlobalIds, globalDisplacement, eleRotationalDisplacement);

            const auto eleTorsionalDisplacementDiff = coeffMatrix * (TEle * eleRotationalDisplacement);
            // 切应变
            const auto gamma = eleTorsionalDisplacementDiff (0, 0) / L;
            // 轴上最大应变和应力
            eleTorsionalStrain = gamma * rodData.m_C;
            eleTorsionalStress = material.m_G * eleTorsionalStrain;
            eleTorque = gamma * rodData.m_J * material.m_G;
        }

        rodStrainSet.emplace_back (RODStrainData{rodData.m_id, eleAxialStrain, 0.0, eleTorsionalStrain, 0.0});
        rodStressSet.emplace_back (RODStressData{rodData.m_id, eleAxialStress, 0.0, eleTorsionalStress, 0.0});
        rodForceSet.emplace_back (RODForceData{rodData.m_id, eleAxialForce, eleTorque});
    }

    if (ElementTypes::ELETYPE_Link == m_type)
    {
        dbServiceFactory->get<IRODStrainService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodStrainSet);
        dbServiceFactory->get<IRODStressService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodStressSet);
        dbServiceFactory->get<IRODForceService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodForceSet);
    }
    else
    {
        dbServiceFactory->get<ICONRODStrainService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodStrainSet);
        dbServiceFactory->get<ICONRODStressService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodStressSet);
        dbServiceFactory->get<ICONRODForceService> ()->AddBySubcaseId (static_cast<Id> (subcaseId), rodForceSet);
    }
}

// 计算全局坐标系下单元应力刚度矩阵
void RodCalculator::ComputeStressStiffness (const SG::DBManager::DBServiceFactorySP dbServiceFactory, int staticSubId)
{
    std::vector<RODForceData>    rodForceSet;
    std::vector<CONRODForceData> conrodForceSet;
    dbServiceFactory->get<IRODForceService> ()->Find ({SOLType::SOL_SESTATIC ,static_cast<Id> (staticSubId)} , _OUT rodForceSet);
    dbServiceFactory->get<ICONRODForceService> ()->Find ({SOLType::SOL_SESTATIC,static_cast<Id> (staticSubId)}, _OUT conrodForceSet);

    auto&                    assembly       = AssemblyEigen::GetInstance ();
    const auto&              coordProcessor = CoordinateProcessor::GetInstance ();
    CONRODData               rodData;
    std::vector<std::size_t> nodeIndexes (2);
    std::vector<int>         globalIndex (12);
    IdArray                  dispCoordIds (2);
    PointArray               points (2);
    Matrixd                  StressStiffness (6, 6);

    static const DofSet DOFs{DOFDirections::DOF_Ux, DOFDirections::DOF_Uy, DOFDirections::DOF_Uz};

    const auto& nodeService = dbServiceFactory->get<INodeService> ();
    const auto eleSize = dbServiceFactory->get<IElementService>()->GetElementSize (m_type);
    for (std::size_t index{0}; index < eleSize; ++index)
    {
        // 根据不同类型，获取单元数据
        GetElement (dbServiceFactory, m_type, index, _OUT rodData);
        ModelInfoQuery::QueryNode (nodeService, rodData.m_g, 2, points, nodeIndexes, dispCoordIds);

        auto xVec = SG::Algebra::Distance (points[1], points[0]);
        Real L;
        xVec.Normalize (L);
        const auto& TEle = CoordTrans (xVec, coordProcessor, dispCoordIds, points);

        Real FxL;
        if (ElementTypes::ELETYPE_Link == m_type)
        {
            FxL = rodForceSet[index].m_AF / L;
        }
        else
        {
            FxL = conrodForceSet[index].m_AF / L;
        }

        const Matrixd StressStiffness1 (2, 2, {FxL, -FxL, -FxL, FxL});
        const Matrixd StressStiffness = TEle.Transpose () * StressStiffness1 * TEle;

        assembly.GetGlobalDOFId (nodeIndexes, DOFs, _OUT globalIndex);
        assembly.AccumulateStressStiffness (globalIndex, StressStiffness);
    }
}